import logging
import os
import time
import torch
import torch.nn as nn
from utils.meter import AverageMeter
from utils.metrics import R1_mAP_eval
from torch.amp import autocast, GradScaler
import torch.distributed as dist

def do_train(cfg,
             model,
             center_criterion,
             train_loader,
             val_loader,
             optimizer,
             optimizer_center,
             scheduler,
             loss_fn,
             num_query, local_rank):
    log_period = cfg.SOLVER.LOG_PERIOD
    checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD
    eval_period = cfg.SOLVER.EVAL_PERIOD
    output_dir = cfg.OUTPUT_DIR
    device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
    epochs = cfg.SOLVER.MAX_EPOCHS

    logger = logging.getLogger("transreid.train")
    logger.info('start training')
    logger.info(f"Device: {device}, CUDA available: {torch.cuda.is_available()}")
    _LOCAL_PROCESS_GROUP = None

    model.to(device)
    logger.info("Model moved to device")
    if torch.cuda.device_count() > 1 and cfg.MODEL.DIST_TRAIN:
        print(f'Using {torch.cuda.device_count()} GPUs for training')
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], find_unused_parameters=True)
        logger.info("Initialized DistributedDataParallel")

    loss_meter = AverageMeter()
    acc_meter = AverageMeter()

    evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
    scaler = GradScaler('cuda' if torch.cuda.is_available() else 'cpu')
    logger.info("GradScaler initialized")

    for epoch in range(1, epochs + 1):
        start_time = time.time()
        loss_meter.reset()
        acc_meter.reset()
        evaluator.reset()
        scheduler.step(epoch)
        model.train()
        logger.info(f"Epoch {epoch} started")

        for n_iter, (img, vid, target_cam, target_view) in enumerate(train_loader):
            logger.info(f"Iteration {n_iter+1}: Data loaded")
            optimizer.zero_grad()
            if center_criterion is not None and optimizer_center is not None:
                optimizer_center.zero_grad()

            img = img.to(device)
            target = vid.to(device)
            target_cam = target_cam.to(device)
            target_view = target_view.to(device)
            logger.info(f"Iteration {n_iter+1}: Data moved to {device}")

            with autocast('cuda' if torch.cuda.is_available() else 'cpu'):
                score, feat = model(img, target, cam_label=target_cam, view_label=target_view)
                logger.info(f"Iteration {n_iter+1}: Model forward pass completed")
                loss = loss_fn(score, feat, target, target_cam)
                logger.info(f"Iteration {n_iter+1}: Loss computed")

            scaler.scale(loss).backward()
            logger.info(f"Iteration {n_iter+1}: Backward pass completed")
            scaler.step(optimizer)
            scaler.update()
            logger.info(f"Iteration {n_iter+1}: Optimizer step completed")

            if isinstance(score, list):
                acc = (score[0].max(1)[1] == target).float().mean()
            else:
                acc = (score.max(1)[1] == target).float().mean()

            loss_meter.update(loss.item(), img.shape[0])
            acc_meter.update(acc, 1)

            if torch.cuda.is_available():
                torch.cuda.synchronize()
            if (n_iter + 1) % log_period == 0:
                logger.info(f"Epoch[{epoch}] Iteration[{n_iter+1}/{len(train_loader)}] Loss: {loss_meter.avg:.3f}, Acc: {acc_meter.avg:.3f}, Base Lr: {scheduler._get_lr(epoch)[0]:.2e}")

        end_time = time.time()
        time_per_batch = (end_time - start_time) / (n_iter + 1)
        logger.info(f"Epoch {epoch} done. Time per batch: {time_per_batch:.3f}[s] Speed: {train_loader.batch_size / time_per_batch:.1f}[samples/s]")

        if epoch % checkpoint_period == 0:
            if cfg.MODEL.DIST_TRAIN:
                if dist.get_rank() == 0:
                    torch.save(model.state_dict(),
                               os.path.join(output_dir, f"{cfg.MODEL.NAME}_{epoch}.pth"))
            else:
                torch.save(model.state_dict(),
                           os.path.join(output_dir, f"{cfg.MODEL.NAME}_{epoch}.pth"))

        if epoch % eval_period == 0:
            model.eval()
            for n_iter, (img, vid, camid, camids, target_view, _) in enumerate(val_loader):
                with torch.no_grad():
                    img = img.to(device)
                    camids = camids.to(device)
                    target_view = target_view.to(device)
                    feat = model(img, cam_label=camids, view_label=target_view)
                    evaluator.update((feat, vid, camid))
            cmc, mAP, _, _, _, _, _ = evaluator.compute()
            logger.info(f"Validation Results - Epoch: {epoch}")
            logger.info(f"mAP: {mAP:.1%}")
            for r in [1, 5, 10]:
                logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.1%}")
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

    logger.info("Training finished")

def do_inference(cfg,
                model,
                val_loader,
                num_query):
    device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
    logger = logging.getLogger("transreid.test")
    logger.info("Enter inferencing")

    evaluator = R1_mAP_eval(num_query, max_rank=50, feat_norm=cfg.TEST.FEAT_NORM)
    evaluator.reset()

    model.to(device)
    if torch.cuda.device_count() > 1:
        print(f'Using {torch.cuda.device_count()} GPUs for inference')
        model = nn.DataParallel(model)

    model.eval()
    img_path_list = []

    for n_iter, (img, pid, camid, camids, target_view, imgpath) in enumerate(val_loader):
        with torch.no_grad():
            img = img.to(device)
            camids = camids.to(device)
            target_view = target_view.to(device)
            feat = model(img, cam_label=camids, view_label=target_view)
            evaluator.update((feat, pid, camid))
            img_path_list.extend(imgpath)

    cmc, mAP, _, _, _, _, _ = evaluator.compute()
    logger.info("Validation Results ")
    logger.info(f"mAP: {mAP:.1%}")
    for r in [1, 5, 10]:
        logger.info(f"CMC curve, Rank-{r:<3}:{cmc[r - 1]:.1%}")
    return cmc[0], cmc[4]